from pdf_craft import PDFPageExtractor, MarkDownWriter,ExtractedTableFormat
import os
from fastapi import APIRouter, UploadFile, File, HTTPException, Query
from typing import Optional

# 创建路由器
router = APIRouter(
    prefix="/pdf2md",
    tags=["PDF转Markdown"],
    responses={404: {"description": "Not found"}},
)

# 默认目录
DEFAULT_INPUT_DIR = "input_files"
DEFAULT_OUTPUT_DIR = "output_files"

textextractor = PDFPageExtractor(
  device="cpu", # If you want to use CUDA, please change to device="cuda" format.
  model_dir_path="modelDir", # The folder address where the AI ​​model is downloaded and installed
)

# 解析表格必须要cuda-CUDA is not available. You cannot parse table from image.
imagePdfExtractor = PDFPageExtractor(
  device="cuda", 
  model_dir_path="modelDir", 
  ocr_level="OncePerBlock",
  extract_formula=True, # Enable formula recognition
  extract_table_format=ExtractedTableFormat.MARKDOWN
)

def textPdf2md(pdf_path,markdown_path):
  
  with MarkDownWriter(markdown_path, "images", "utf-8") as md:
    for block in textextractor.extract(pdf=pdf_path):
      md.write(block)
      
      
def imagePdf2md(pdf_path,markdown_path):
  with MarkDownWriter(markdown_path, "images", "utf-8") as md:
    for block in imagePdfExtractor.extract(pdf=pdf_path):
      md.write(block)
      
# 用于API的工具函数
def ensure_directory_exists(directory_path):
    """确保目录存在，如果不存在则创建"""
    # 检查路径是否为空字符串或None
    if not directory_path:
        return  # 如果是空字符串或None，直接返回，不创建目录
        
    if not os.path.exists(directory_path):
        os.makedirs(directory_path)

# 用于API的函数版本
def process_text_pdf_to_md(pdf_path, output_dir="output_files"):
    """
    处理文本PDF并返回生成的Markdown文件路径
    
    Args:
        pdf_path: PDF文件路径
        output_dir: 输出目录，默认为'output_files'
        
    Returns:
        生成的Markdown文件路径
    """
    # 确保输出目录存在
    ensure_directory_exists(output_dir)
    
    # 生成唯一的输出文件名
    output_filename = f"output_text_{os.path.basename(pdf_path).replace('.pdf', '')}.md"
    markdown_path = os.path.join(output_dir, output_filename)
    
    # 调用原始函数进行转换
    textPdf2md(pdf_path, markdown_path)
    
    return markdown_path

def process_image_pdf_to_md(pdf_path, output_dir="output_files"):
    """
    处理图片PDF并返回生成的Markdown文件路径
    
    Args:
        pdf_path: PDF文件路径
        output_dir: 输出目录，默认为'output_files'
        
    Returns:
        生成的Markdown文件路径
    """
    # 确保输出目录存在
    ensure_directory_exists(output_dir)
    
    # 生成唯一的输出文件名
    output_filename = f"output_image_{os.path.basename(pdf_path).replace('.pdf', '')}.md"
    markdown_path = os.path.join(output_dir, output_filename)
    
    # 调用原始函数进行转换
    imagePdf2md(pdf_path, markdown_path)
    
    return markdown_path

# 添加路由器端点
@router.post("/text")
async def convert_text_pdf_to_md(
    file: UploadFile = File(...),
    input_dir: Optional[str] = Query(DEFAULT_INPUT_DIR, description="输入文件目录"),
    output_dir: Optional[str] = Query(DEFAULT_OUTPUT_DIR, description="输出文件目录")
):
    """
    将文本类型的PDF文件转换为Markdown
    
    - **file**: 要上传的PDF文件
    - **input_dir**: 输入文件保存目录，默认为'input_files'
    - **output_dir**: 输出文件保存目录，默认为'output_files'
    """
    # 处理默认目录
    if not input_dir or input_dir.strip() == "":
        input_dir = DEFAULT_INPUT_DIR
    
    if not output_dir or output_dir.strip() == "":
        output_dir = DEFAULT_OUTPUT_DIR
    
    # 确保目录存在
    ensure_directory_exists(input_dir)
    ensure_directory_exists(output_dir)
    
    # 检查文件类型
    if not file.filename.lower().endswith('.pdf'):
        raise HTTPException(status_code=400, detail="只接受PDF文件")
    
    # 保存上传的文件
    file_path = os.path.join(input_dir, file.filename)
    with open(file_path, "wb") as buffer:
        buffer.write(await file.read())
    
    try:
        # 处理PDF文件，传入输出目录
        output_path = process_text_pdf_to_md(file_path, output_dir=output_dir)
        
        # 返回文件下载URL（修正路径处理）
        file_name = os.path.basename(output_path)
        download_url = f"/outputs/{file_name}"
        
        return {
            "success": True,
            "message": "文本PDF转换成功",
            "file_path": output_path,
            "download_url": download_url
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"处理文件时出错: {str(e)}")

@router.post("/image")
async def convert_image_pdf_to_md(
    file: UploadFile = File(...),
    input_dir: Optional[str] = Query(DEFAULT_INPUT_DIR, description="输入文件目录"),
    output_dir: Optional[str] = Query(DEFAULT_OUTPUT_DIR, description="输出文件目录")
):
    """
    将图像类型的PDF文件转换为Markdown
    
    - **file**: 要上传的PDF文件
    - **input_dir**: 输入文件保存目录，默认为'input_files'
    - **output_dir**: 输出文件保存目录，默认为'output_files'
    """
    # 处理默认目录
    if not input_dir or input_dir.strip() == "":
        input_dir = DEFAULT_INPUT_DIR
    
    if not output_dir or output_dir.strip() == "":
        output_dir = DEFAULT_OUTPUT_DIR
    
    # 确保目录存在
    ensure_directory_exists(input_dir)
    ensure_directory_exists(output_dir)
    
    # 检查文件类型
    if not file.filename.lower().endswith('.pdf'):
        raise HTTPException(status_code=400, detail="只接受PDF文件")
    
    # 保存上传的文件
    file_path = os.path.join(input_dir, file.filename)
    with open(file_path, "wb") as buffer:
        buffer.write(await file.read())
    
    try:
        # 处理PDF文件，传入输出目录
        output_path = process_image_pdf_to_md(file_path, output_dir=output_dir)
        
        # 返回文件下载URL（修正路径处理）
        file_name = os.path.basename(output_path)
        download_url = f"/outputs/{file_name}"
        
        return {
            "success": True,
            "message": "图像PDF转换成功",
            "file_path": output_path,
            "download_url": download_url
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"处理文件时出错: {str(e)}")
      
